#include <iostream>
#include <cassert>
#include "../OpenclObj.hpp"
#include "../matrix.hpp"
using namespace std;

int main()
{

    OpenclObj clobj;
    clobj.GetPlatformID(1);
    clobj.CreateContextFromType(CL_DEVICE_TYPE_CPU);
    clobj.CreateCommandQueue();
    // 默认开启 cl1.2版本， 归约函数为 cl2.0
    string options{"-cl-std=CL2.0 -cl-opt-disable -Werror"};
    clobj.CreateProgram("vec_add.cl", options.c_str());
    // clobj.CreateProgram("vec_add.cl", NULL);
    clobj.CreateKernelsInProgram();

    // 归约中 一个工作组中的工作项数目不超过 4096 超过就不正确了
    constexpr size_t NUM = 4096;
    Matrix<double> a(1, NUM), b(1, NUM), c(1, NUM);
    // a.InitRandom(1.1, 9.1);
    // b.InitRandom(1.1, 9.1);
    a.initElem();
    b.initElem();

    cl_int err;
    cl_mem mem_a, mem_b, mem_c;
    mem_a = clCreateBuffer(clobj.GetContext(),
                           CL_MEM_COPY_HOST_PTR | CL_MEM_READ_ONLY,
                           sizeof(double) * NUM, a.data(), &err);
    assert(err == CL_SUCCESS);
    // write 是为了下面要用 
    mem_b = clCreateBuffer(clobj.GetContext(),
                           CL_MEM_COPY_HOST_PTR | CL_MEM_READ_WRITE,
                           sizeof(double) * NUM, b.data(), &err);
    assert(err == CL_SUCCESS);
    mem_c = clCreateBuffer(clobj.GetContext(),
                           CL_MEM_READ_WRITE,
                           sizeof(double) * NUM, NULL, &err);
    assert(err == CL_SUCCESS);

    // 两数组相加
    cl_kernel vec_add_ker = clobj.GetKernel("vec_add");
    err = clSetKernelArg(vec_add_ker, 0, sizeof(cl_mem), &mem_a);
    err |= clSetKernelArg(vec_add_ker, 1, sizeof(cl_mem), &mem_b);
    err |= clSetKernelArg(vec_add_ker, 2, sizeof(cl_mem), &mem_c);
    assert(err == CL_SUCCESS);

    cl_uint work_dim = 1;
    size_t global_size[1] = {NUM};
    // 最好选择 2 的指数 
    size_t local_size[1] = {16};
    err = clEnqueueNDRangeKernel(clobj.GetCommandQueue(), vec_add_ker,
                                 work_dim, NULL, global_size,
                                 local_size, 0, NULL, NULL);
    assert(err == CL_SUCCESS);

    clFinish(clobj.GetCommandQueue());
    err = clEnqueueReadBuffer(clobj.GetCommandQueue(), mem_c,
                              CL_TRUE, 0, sizeof(double) * NUM,
                              c.data(), 0, NULL, NULL);
    assert(err == CL_SUCCESS);
    Matrix<double> d = a + b;
    d.print();
    c.print();
    if (d == c)
    {
        cout << "equal\n";
    }
    else
    {
        cout << "not equal\n";
    }

    // 归约测试  归约中的 local_size 不能太大，否则错误
    // a.print();
    double sum{0.0};
    for (size_t i = 0; i < NUM; ++i)
        sum += a(0, i);
    double max{0.0};
    for (size_t i = 0; i < NUM; ++i)
        max = (a(0, i) < max) ? max : a(0, i);
    double min{0.0};
    for (size_t i = 0; i < NUM; ++i)
        min = (a(0, i) < min) ? a(0, i) : min;

    cl_kernel vec_sum_ker = clobj.GetKernel("vec_sum");
    err = clSetKernelArg(vec_sum_ker, 0, sizeof(cl_mem), &mem_a);
    err |= clSetKernelArg(vec_sum_ker, 1, sizeof(cl_mem), &mem_b);
    assert(err == CL_SUCCESS);

    work_dim = 1;
    global_size[0] = NUM;
    local_size[0] = NUM;
    // 复制黏贴的时候要注意修改
    err = clEnqueueNDRangeKernel(clobj.GetCommandQueue(), vec_sum_ker,
                                 work_dim, NULL, global_size,
                                 local_size, 0, NULL, NULL);

    clFinish(clobj.GetCommandQueue());
    err = clEnqueueReadBuffer(clobj.GetCommandQueue(), mem_b,
                              CL_TRUE, 0, sizeof(double) * NUM,
                              b.data(), 0, NULL, NULL);
    assert(err == CL_SUCCESS);
    // b.print();

    // 矩阵下标是从 0 开始， 不是 1
    cout << " serial sum = " << sum << "\n";
    // cout << " opencl sum = " << b(1, 0) << "\n";
    cout << " opencl reduce_add = " << b(0, 1) << "\n";

    cout << " serial min = " << min << "\n";
    cout << " opencl reduce min = " << b(0, 2) << "\n";

    cout << " serial max = " << max << "\n";
    cout << " opencl reduce max = " << b(0, 3) << "\n";
    
    clReleaseMemObject(mem_a);
    clReleaseMemObject(mem_b);
    clReleaseMemObject(mem_c);
    return 0;
}